
#include <mach/mach_types.h>
#include <mach-o/loader.h>
#include <mach-o/fat.h>

#include <sys/systm.h>
#include <sys/kauth.h>
#include <sys/vnode.h>
#include <sys/proc.h>
#include <sys/file.h>
#include <sys/socket.h>
#include <sys/kern_control.h>
#include <sys/ioctl.h>
#include <sys/sys_domain.h>

#include <libkern/OSMalloc.h>

#include "MacFireWall.h"

kern_return_t Drv_start(kmod_info_t * ki, void *d);
kern_return_t Drv_stop(kmod_info_t *ki, void *d);

int getLogLength();
int moveOffset();
void writeToLog(char * text);

static int offset = 0;
static int isWriteToLogFirst = 1;
static char previousChar = '\0';
struct list *l;

/* SOCKET PART */

char g_string_buf[1024]={0};
static kern_ctl_ref g_ctl_ref;
static boolean_t g_filter_registered = FALSE;

static int hello_ctl_set(kern_ctl_ref ctl_ref, u_int32_t unit, void* unitinfo, int opt, void* data, size_t len);
static int hello_ctl_connect(kern_ctl_ref ctl_ref, struct sockaddr_ctl *sac, void** unitinfo);
static errno_t hello_ctl_disconnect(kern_ctl_ref ctl_ref, u_int32_t unit, void* unitinfo);

static struct kern_ctl_reg g_kern_ctl_reg = {
    "com.osxkernel.HelloKernControl",
    0,
    0,
    CTL_FLAG_PRIVILEGED,
    0,
    0,
    hello_ctl_connect,
    hello_ctl_disconnect,
    NULL,
    hello_ctl_set,
    NULL
};



/* LISTEN-BLOCK ;  READ-WRITE MAGIC KEY; */

//void writeMagicKey(char path[], uint16_t magicKey[]); //write Mkey to block list


void checkMagicKey(char *path,vnode_t vn,vfs_context_t vfs,uint8_t magicKey[]); // get MKey from exec vnode

static kauth_listener_t processListener;

static int callbackProcessListener  // process Listener/Blocker
(
 kauth_cred_t    credential,
 void *          idata,
 kauth_action_t  action,
 uintptr_t       arg0,
 uintptr_t       arg1,
 uintptr_t       arg2,
 uintptr_t       arg3
 );

void getExPath(char path[]){
    
    int p = strlen(path);
    
    for(int i = p; i>0;i--)
        if(path[i]=='/'){
            p=++i;
            break;
        }

    char ex[20] = {0};
    for(int i = 0; i<20;i++){
        if(path[p]=='.') break;
        ex[i] = path[p];
        p++;
    }
     strncat(path, "/Contents/MacOS/", strlen("/Contents/MacOS/"));
     strncat(path, ex, strlen(ex));
}


kern_return_t Drv_start(kmod_info_t * ki, void *d)
{
    l = new_list();
    
    if(getLogLength() == 1)
    {
        offset = 0;
    }

    int ret = ctl_register(&g_kern_ctl_reg, &g_ctl_ref);
    if (ret == KERN_SUCCESS)
    {
        
        processListener = kauth_listen_scope(KAUTH_SCOPE_VNODE, callbackProcessListener,NULL);
        printf("Start\n");
        writeToLog("FireWall been started\r\n");
        g_filter_registered = TRUE;
        
        return KERN_SUCCESS;
    }
    
    else return KERN_FAILURE;
}

kern_return_t Drv_stop(kmod_info_t *ki, void *d)

{   
   kauth_unlisten_scope(processListener);

   processListener = NULL;
    
   delete_list(l);
    
    if (g_filter_registered)
        ctl_deregister(g_ctl_ref);
    
    printf("Stop\n");
    writeToLog("FireWall been stopped\r\n");
    return KERN_SUCCESS;
}


/* SOCKET PART */

static int hello_ctl_connect(kern_ctl_ref ctl_ref, struct sockaddr_ctl *sac, void** unitinfo)
{
   // printf("process with pid=%d connected\n", proc_selfpid());
    return 0;
}

static errno_t hello_ctl_disconnect(kern_ctl_ref ctl_ref, u_int32_t unit, void* unitinfo)
{
    //printf("process with pid=%d disconnected\n", proc_selfpid());
    return 0;
}

static int hello_ctl_set(kern_ctl_ref ctl_ref, u_int32_t unit, void* unitinfo, int opt, void* data, size_t len)
{
    int ret = 0;
    for (int i=0;i<255;i++)
      g_string_buf[i]='\0';
    
    switch (opt) {
            
        case 2:
            strncpy(g_string_buf, (char*)data, min(255, len));
            printf("Kext get a path: %s\n", g_string_buf);
            writeToLog("ADD BLOCKED: ");
            writeToLog(g_string_buf);
            writeToLog("\r\n");
            
            getExPath(g_string_buf);
            insert_item_uniq(l, g_string_buf);
            break;
            
        case 3:
            strncpy(g_string_buf, (char*)data, min(255, len));
            printf("Kext delete MK for path: %s\n", g_string_buf);
            writeToLog("DELETE BLOCKED: ");
            writeToLog(g_string_buf);
            writeToLog("\r\n");
            getExPath(g_string_buf);
            delete_item(l,g_string_buf);
            break;
            
        default:
            ret = ENOTSUP;
        break;
    }
    return ret;
}

/* Listen/ Read MK/ Write MK  PART */

static int callbackProcessListener
(
 kauth_cred_t    credential,
 void *          idata,
 kauth_action_t  action,
 uintptr_t       arg0,
 uintptr_t       arg1,
 uintptr_t       arg2,
 uintptr_t       arg3
 )
{
      
    char path[1024]= {0};
    uint8_t mKey[16]={0};
    int size = 1024;
    if(KAUTH_VNODE_EXECUTE == action){
        vn_getpath((vnode_t) arg1,path,&size);
         
        if(vnode_vtype((vnode_t)arg1)==VREG){
    
            checkMagicKey(path, (vnode_t)arg1, (vfs_context_t)arg0, mKey);
        
         
            if(find_item(l, mKey))
            {
                //printf("blocked\n");
                writeToLog("Trying start ");
                writeToLog(path);
                writeToLog("\r\nPERMISSION DENIED\r\n");
              return  KAUTH_RESULT_DENY;
            
            }
        }

    }
    
    return KAUTH_RESULT_DEFER;
}

void checkMagicKey(char *path,vnode_t vn,vfs_context_t vfs,uint8_t magicKey[]){
   
    char adr[4096]= {0};
    struct uuid_command *uid = NULL;
    
    if(strlen(path)>2){
        
        if (vnode_open(path, O_RDONLY, (0), 0, &vn, vfs) == 0) {
            uio_t read;
            read = uio_create(1, 0, UIO_SYSSPACE, UIO_READ);
            uio_addiov(read, (user_addr_t)adr, 4096);
            
            VNOP_READ(vn, read,  0, vfs);
                      
            uid = (struct uuid_command*)adr;
            for(int i =0; i<16;i++){
                magicKey[i]=uid->uuid[i];
            }
        
            uio_free(read);
            vnode_close(vn, 0, vfs);
        } else {
          
        }
    }
}

/* LOG PART */

int getLogLength()
{
    int result = 0;
    while(result == 0)
    {
        result = moveOffset();
    }
    return result;
}

int moveOffset()
{
    vnode_t vn;
    vfs_context_t vfs = vfs_context_current();
    int result = vnode_open("/var/log/firewall.log", FREAD, (0), 0, &vn, vfs);
    if (result == 0)
    {
        char buf[1];
        uio_t read;
        read = uio_create(1, offset, UIO_SYSSPACE, UIO_READ);
        uio_addiov(read, (user_addr_t)buf, 1);
        int resultRead = VNOP_READ(vn, read, 0, vfs);
        uio_free(read);
        vnode_close(vn, 0, vfs);
        if(resultRead == 0 && buf[0] != '\0' && (buf[0] != '\n' || buf[0] != previousChar))
        {
            offset++;
          //  printf("new char is:%d\n", buf[0]);
            previousChar = buf[0];
            return 0;
        }
      //  printf("Reading log stopped at %d\n", offset);
        return 2;
    }
  //  printf("Can't open log\n");
    return 1;
}

void writeToLog(char text[])
{
    vnode_t vn;
    vfs_context_t vfs = vfs_context_current();
    int result = 1;
    if(isWriteToLogFirst == 1)
    {
        result = vnode_open("/var/log/firewall.log",( FWRITE | O_CREAT | O_TRUNC ), (0), 0, &vn, vfs);
       
    }
    else
    {
        result = vnode_open("/var/log/firewall.log",( FWRITE | O_CREAT | O_APPEND ), (0), 0, &vn, vfs);
    }
    if (result == 0)
    {
        uio_t write;
        write = uio_create(1, offset, UIO_SYSSPACE, UIO_WRITE);
        uio_addiov(write, (user_addr_t)text, strlen(text));
        VNOP_WRITE(vn, write, 0, vfs);
        uio_free(write);
        vnode_close(vn, 0, vfs);
        offset += strlen(text);
        isWriteToLogFirst = 0;
    }
    else
    {
        printf("Can't write to log\n");
    }
}


